package data_store

import (
	"sync"

	"atomgit.com/openlinksaas/proto-gen-go.git/trace_proxy_api"
	"github.com/edwingeng/deque/v2"
	"go.opentelemetry.io/collector/pdata/pcommon"
	"go.opentelemetry.io/collector/pdata/ptrace"
)

type EventRecord struct {
	Name    string
	Time    int64
	AttrMap pcommon.Map
}

type SpanRecord struct {
	SpanId         string
	ParentSpanId   string
	TraceId        string
	ServiceName    string
	ServiceVersion string
	SpanName       string
	StartTime      int64
	EndTime        int64
	SpanKind       ptrace.SpanKind
	AttrMap        pcommon.Map
	EventList      []*EventRecord
}

func (record *SpanRecord) ToSpanInfo() *trace_proxy_api.SpanInfo {
	attrList := []*trace_proxy_api.AttrInfo{}
	record.AttrMap.Range(func(key string, value pcommon.Value) bool {
		attrList = append(attrList, &trace_proxy_api.AttrInfo{
			Key:   key,
			Value: value.AsString(),
		})
		return true
	})
	eventList := []*trace_proxy_api.EventInfo{}
	for _, ev := range record.EventList {
		evAttrList := []*trace_proxy_api.AttrInfo{}
		ev.AttrMap.Range(func(key string, value pcommon.Value) bool {
			evAttrList = append(evAttrList, &trace_proxy_api.AttrInfo{
				Key:   key,
				Value: value.AsString(),
			})
			return true
		})
		eventList = append(eventList, &trace_proxy_api.EventInfo{
			Name:      ev.Name,
			TimeStamp: ev.Time,
			AttrList:  evAttrList,
		})
	}
	return &trace_proxy_api.SpanInfo{
		SpanId:         record.SpanId,
		ParentSpanId:   record.ParentSpanId,
		TraceId:        record.TraceId,
		ServiceName:    record.ServiceName,
		ServiceVersion: record.ServiceVersion,
		SpanName:       record.SpanName,
		StartTimeStamp: record.StartTime,
		EndTimeStamp:   record.EndTime,
		SpanKind:       trace_proxy_api.SPAN_KIND(record.SpanKind),
		AttrList:       attrList,
		EventList:      eventList,
	}
}

type SpanTable struct {
	spanMap    map[string]*SpanRecord
	traceIndex map[string]*deque.Deque[string]
	mutex      sync.Mutex
}

func newSpanTable(maxTrace uint32) *SpanTable {
	return &SpanTable{
		spanMap:    make(map[string]*SpanRecord, maxTrace*6),
		traceIndex: make(map[string]*deque.Deque[string], maxTrace),
	}
}

func (tab *SpanTable) Insert(record *SpanRecord) {
	tab.mutex.Lock()
	defer tab.mutex.Unlock()

	_, ok := tab.spanMap[record.SpanId]
	if ok {
		return
	}

	tab.spanMap[record.SpanId] = record
	traceQueue, ok := tab.traceIndex[record.TraceId]
	if !ok {
		traceQueue = deque.NewDeque[string]()
		tab.traceIndex[record.TraceId] = traceQueue
	}
	traceQueue.PushBack(record.SpanId)
}

func (tab *SpanTable) RemoveByTrace(traceId string) {
	tab.mutex.Lock()
	defer tab.mutex.Unlock()

	traceQueue, ok := tab.traceIndex[traceId]
	if !ok {
		return
	}
	traceQueue.Range(func(_ int, spanId string) bool {
		delete(tab.spanMap, spanId)
		return true
	})
	delete(tab.traceIndex, traceId)
}

func (tab *SpanTable) ListByTrace(traceId string) []*SpanRecord {
	tab.mutex.Lock()
	defer tab.mutex.Unlock()

	retList := []*SpanRecord{}

	traceQueue, ok := tab.traceIndex[traceId]
	if !ok {
		return retList
	}

	traceQueue.Range(func(_ int, spanId string) bool {
		record, ok := tab.spanMap[spanId]
		if ok {
			retList = append(retList, record)
		}
		return true
	})

	return retList
}

func (tab *SpanTable) Get(spanId string) *SpanRecord {
	tab.mutex.Lock()
	defer tab.mutex.Unlock()

	return tab.spanMap[spanId]
}
